import os
import numpy as np
import cv2
import matplotlib
# 设置非交互式后端，避免tkinter线程问题
matplotlib.use('Agg')
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
import matplotlib.gridspec as gridspec
from matplotlib.font_manager import FontProperties
import pandas as pd
import uuid
from flask import current_app

def calculate_area(mask_path, latitude=None, longitude=None):
    """
    计算水域面积，并支持用户绘制ROI计算特定区域面积
    
    Args:
        mask_path (str): 掩码图像路径
        latitude (float): 纬度，用于地理坐标转换
        longitude (float): 经度，用于地理坐标转换
    
    Returns:
        dict: 包含面积计算结果的字典
    """
    # 读取掩码
    if os.path.exists(mask_path):
        mask = cv2.imread(mask_path)
        if mask is None:
            raise FileNotFoundError(f"无法读取掩码图像: {mask_path}")
    else:
        raise FileNotFoundError(f"掩码图像不存在: {mask_path}")
    
    # 计算总水域面积（像素数量）
    if len(mask.shape) == 3:
        # 提取蓝色通道的掩码
        water_mask = (mask[:, :, 0] == 0) & (mask[:, :, 1] == 0) & (mask[:, :, 2] > 200)
    else:
        water_mask = mask > 0
    
    total_pixels = np.sum(water_mask)
    image_area = mask.shape[0] * mask.shape[1]
    
    # 根据经纬度信息计算实际面积（假设每个像素代表的实际面积）
    # 注意：这是一个简化的计算方法，实际应用中需要更精确的投影转换
    if latitude is not None and longitude is not None:
        # 根据纬度计算像素对应的实际距离（简化估算）
        # 赤道上1度约为111公里，随着纬度增加，经度间距离减小
        lat_rad = np.radians(abs(latitude))
        meters_per_degree_lat = 111320  # 每度纬度约111.32公里
        meters_per_degree_lon = 111320 * np.cos(lat_rad)  # 每度经度随纬度变化
        
        # 假设影像覆盖的经纬度范围（简化估算）
        # 实际应该从影像元数据中获取
        coverage_degrees_lat = 0.1  # 覆盖的纬度范围
        coverage_degrees_lon = 0.1  # 覆盖的经度范围
        
        # 计算实际面积
        image_width_meters = coverage_degrees_lon * meters_per_degree_lon
        image_height_meters = coverage_degrees_lat * meters_per_degree_lat
        
        # 计算实际面积（平方米）
        pixel_area_sqm = (image_width_meters * image_height_meters) / image_area
        total_area_sqm = total_pixels * pixel_area_sqm
        
        # 转换为平方公里
        total_area_sqkm = total_area_sqm / 1000000
    else:
        # 如果没有地理信息，返回像素比例
        total_area_sqkm = None
        pixel_area_sqm = None
    
    # 计算水域占总面积的百分比
    area_percentage = (total_pixels / image_area) * 100
    
    return {
        'total_pixels': int(total_pixels),
        'image_area': int(image_area),
        'area_percentage': float(area_percentage),
        'total_area_sqkm': float(total_area_sqkm) if total_area_sqkm else None,
        'pixel_area_sqm': float(pixel_area_sqm) if pixel_area_sqm else None
    }

def generate_time_series(model_predictions):
    """
    生成时间序列分析图表
    
    Args:
        model_predictions (dict): 包含各模型预测结果的字典
    
    Returns:
        str: 生成的图表文件名
    """
    # 确保使用非交互式后端
    plt.switch_backend('Agg')
    
    # 设置中文字体
    try:
        font = FontProperties(fname=r"C:\Windows\Fonts\simsun.ttc", size=12)
    except:
        font = FontProperties(size=12)
    
    # 创建图表
    plt.figure(figsize=(12, 8))
    
    # 解析日期和面积数据
    for model_name, predictions in model_predictions.items():
        if not predictions:  # 跳过没有数据的模型
            continue
        
        # 按日期排序
        try:
            # 对于字符串格式的日期进行排序
            if isinstance(predictions[0]['date'], str):
                from datetime import datetime
                # 将字符串日期转换为datetime对象进行排序
                for pred in predictions:
                    if isinstance(pred['date'], str):
                        try:
                            # 尝试解析日期字符串 (格式: YYYY-MM-DD HH:MM:SS)
                            pred['date_obj'] = datetime.strptime(pred['date'], '%Y-%m-%d %H:%M:%S')
                        except ValueError:
                            # 如果解析失败，使用当前时间
                            pred['date_obj'] = datetime.now()
                
                predictions_sorted = sorted(predictions, key=lambda x: x['date_obj'])
                # 使用字符串日期显示
                dates = [pred['date'] for pred in predictions_sorted]
            else:
                # 直接使用日期对象排序
                predictions_sorted = sorted(predictions, key=lambda x: x['date'])
                dates = [pred['date'] for pred in predictions_sorted]
        except (KeyError, IndexError, TypeError):
            # 处理异常情况
            predictions_sorted = predictions
            dates = [f"样本{i+1}" for i in range(len(predictions))]
        
        areas = [pred['water_area'] if pred['water_area'] is not None else 0 for pred in predictions_sorted]
        
        # 绘制折线图（使用数字索引替代日期）
        indices = list(range(len(dates)))
        plt.plot(indices, areas, marker='o', linestyle='-', label=model_name.upper())
        
    # 设置x轴刻度标签
    if 'predictions_sorted' in locals() and len(predictions_sorted) > 0:
        plt.xticks(range(len(dates)), dates, rotation=45)
    
    # 添加图表标题和标签
    plt.title('水域面积随时间变化趋势', fontproperties=font)
    plt.xlabel('日期', fontproperties=font)
    plt.ylabel('水域面积比例', fontproperties=font)
    plt.grid(True)
    plt.legend(prop=font)
    
    # 调整布局
    plt.tight_layout()
    
    try:
        # 保存图表
        chart_filename = f"time_series_{uuid.uuid4().hex}.png"
        chart_path = os.path.join(current_app.config['RESULTS_FOLDER'], chart_filename)
        plt.savefig(chart_path, dpi=300)
        plt.close('all')  # 确保关闭所有图表，防止内存泄漏
        
        return chart_filename
    except Exception as e:
        print(f"保存图表时出错: {e}")
        # 确保关闭图表
        plt.close('all')
        # 返回一个默认图表名
        return "error_chart.png"

def generate_thematic_map(mask_path, latitude=None, longitude=None, city=None):
    """
    生成包含指北针、比例尺、图例的专题地图
    
    Args:
        mask_path (str): 掩码图像路径
        latitude (float): 纬度
        longitude (float): 经度
        city (str): 城市名称
    
    Returns:
        str: 生成的专题图文件名
    """
    # 确保使用非交互式后端
    plt.switch_backend('Agg')
    
    try:
        # 读取掩码
        mask = cv2.imread(mask_path)
        if mask is None:
            raise FileNotFoundError(f"无法读取掩码图像: {mask_path}")
        
        # 设置中文字体
        try:
            font = FontProperties(fname=r"C:\Windows\Fonts\simsun.ttc", size=12)
        except:
            font = FontProperties(size=12)
        
        # 创建图表
        fig = plt.figure(figsize=(12, 10))
        gs = gridspec.GridSpec(2, 2, width_ratios=[4, 1], height_ratios=[4, 1])
        
        # 主地图
        ax_map = plt.subplot(gs[0, 0])
        ax_map.imshow(cv2.cvtColor(mask, cv2.COLOR_BGR2RGB))
        ax_map.set_title(f"水域分布专题图 - {city if city else '未知区域'}", fontproperties=font)
        ax_map.axis('off')
        
        # 添加坐标信息
        if latitude is not None and longitude is not None:
            ax_map.text(0.02, 0.98, f"经度: {longitude:.4f}°, 纬度: {latitude:.4f}°", 
                        transform=ax_map.transAxes, fontproperties=font,
                        verticalalignment='top', bbox=dict(facecolor='white', alpha=0.7))
        
        # 图例
        ax_legend = plt.subplot(gs[0, 1])
        ax_legend.axis('off')
        
        # 添加图例项
        legend_items = [
            {"color": "blue", "label": "水域"},
            {"color": "green", "label": "植被"},
            {"color": "gray", "label": "建筑"},
            {"color": "tan", "label": "裸地"}
        ]
        
        for i, item in enumerate(legend_items):
            ax_legend.add_patch(Rectangle((0.1, 0.8 - i*0.1), 0.2, 0.05, color=item["color"]))
            ax_legend.text(0.35, 0.8 - i*0.1, item["label"], fontproperties=font, va='center')
        
        ax_legend.set_title("图例", fontproperties=font)
        
        # 指北针
        ax_north = plt.subplot(gs[1, 1])
        ax_north.axis('off')
        
        # 绘制指北针
        circle = plt.Circle((0.5, 0.5), 0.4, fill=False)
        arrow = plt.Arrow(0.5, 0.1, 0, 0.8, width=0.2, color='black')
        ax_north.add_patch(circle)
        ax_north.add_patch(arrow)
        ax_north.text(0.5, 0.9, 'N', ha='center', fontproperties=font)
        ax_north.text(0.9, 0.5, 'E', ha='center', fontproperties=font)
        ax_north.text(0.5, 0.1, 'S', ha='center', fontproperties=font)
        ax_north.text(0.1, 0.5, 'W', ha='center', fontproperties=font)
        
        # 比例尺
        ax_scale = plt.subplot(gs[1, 0])
        ax_scale.axis('off')
        
        # 绘制简单的比例尺
        scale_length = 0.3  # 比例尺长度
        scale_x = 0.1  # 起始x位置
        scale_y = 0.5  # y位置
        
        # 假设的比例尺大小（可以根据实际地理信息调整）
        scale_size = "1 km"
        if latitude is not None:
            # 根据纬度调整比例尺
            scale_size = "1 km" if abs(latitude) < 60 else "500 m"
        
        # 绘制比例尺线
        ax_scale.plot([scale_x, scale_x + scale_length], [scale_y, scale_y], 'k-', lw=2)
        
        # 添加刻度
        for i in range(5):
            x = scale_x + i * scale_length / 4
            ax_scale.plot([x, x], [scale_y - 0.05, scale_y + 0.05], 'k-', lw=1)
        
        # 添加标签
        ax_scale.text(scale_x + scale_length / 2, scale_y - 0.15, scale_size, ha='center', fontproperties=font)
        ax_scale.text(scale_x, scale_y - 0.15, "0", ha='center', fontproperties=font)
        
        # 布局调整
        plt.tight_layout()
        
        # 保存专题图
        map_filename = f"thematic_map_{uuid.uuid4().hex}.png"
        map_path = os.path.join(current_app.config['RESULTS_FOLDER'], map_filename)
        plt.savefig(map_path, dpi=300, bbox_inches='tight')
        plt.close('all')
        
        return map_filename
        
    except Exception as e:
        print(f"生成专题图时出错: {e}")
        plt.close('all')
        return "error_map.png"

def export_to_excel(mask_path, image, prediction):
    """
    导出水域分析结果到Excel文件
    
    Args:
        mask_path (str): 掩码图像路径
        image (models.Image): 图像数据库对象
        prediction (models.Prediction): 预测结果数据库对象
    
    Returns:
        str: 生成的Excel文件路径
    """
    # 读取掩码
    mask = cv2.imread(mask_path)
    if mask is None:
        raise FileNotFoundError(f"无法读取掩码图像: {mask_path}")
    
    # 计算水域面积
    area_data = calculate_area(mask_path, image.latitude, image.longitude)
    
    # 创建DataFrame
    data = {
        '项目': [
            '图像ID', 
            '上传日期', 
            '经度', 
            '纬度', 
            '城市', 
            '模型', 
            '预测日期', 
            '水域像素数', 
            '总像素数', 
            '水域占比(%)', 
            '估计水域面积(km²)', 
            'IoU分数', 
            'F1分数'
        ],
        '值': [
            image.id,
            image.upload_date.strftime('%Y-%m-%d %H:%M:%S'),
            image.longitude,
            image.latitude,
            image.city,
            prediction.model_name,
            prediction.prediction_date.strftime('%Y-%m-%d %H:%M:%S'),
            area_data['total_pixels'],
            area_data['image_area'],
            area_data['area_percentage'],
            area_data['total_area_sqkm'],
            prediction.iou_score,
            prediction.f1_score
        ]
    }
    
    df = pd.DataFrame(data)
    
    # 创建Excel文件
    excel_filename = f"water_analysis_{prediction.id}_{uuid.uuid4().hex}.xlsx"
    excel_path = os.path.join(current_app.config['RESULTS_FOLDER'], excel_filename)
    
    # 创建ExcelWriter对象
    writer = pd.ExcelWriter(excel_path, engine='xlsxwriter')
    
    # 将DataFrame写入Excel
    df.to_excel(writer, sheet_name='水域分析结果', index=False)
    
    # 获取工作簿和工作表对象
    workbook = writer.book
    worksheet = writer.sheets['水域分析结果']
    
    # 设置列宽
    worksheet.set_column('A:A', 18)
    worksheet.set_column('B:B', 25)
    
    # 保存Excel文件
    writer.save()
    
    return excel_path 